import vtk
import myVTKRenderWindowInteractorStyle
reload(myVTKRenderWindowInteractorStyle)


def InitializeVolume():
    """
    Setup a volume  for rendering.

    This has to be called before InitializeRegistration, since the DRR of
    the volume will be the "moving image" in the registration.
    """
    #Create the volume and volume reader
    #This might be a problem if we want to register a 3D object other than a volume (eg. polygonal mesh)
    filename = 'C:/Lib/jointtrack/branches/TestData/vert1.vtk'
    volume = vtk.vtkVolume()
    volumeReader = vtk.vtkStructuredPointsReader()
    volumeReader.SetFileName(filename)
    volumeReader.Update()

    #TODO:Allow the user to modify/customize the transfer functions to change how the DRR is rendered!!
    #Default settings for the opacity transfer function,
    opacityTransferFunction = vtk.vtkPiecewiseFunction()
    opacityTransferFunction.AddPoint(0.0, 0.0)
    opacityTransferFunction.AddPoint(70, 0.0)
    opacityTransferFunction.AddPoint(255, 0.5)

    # Create transfer mapping scalar value to color
    colorTransferFunction = vtk.vtkColorTransferFunction()
    colorTransferFunction.AddRGBPoint(0.0, 0.0, 0.0, 0.0)
    colorTransferFunction.AddRGBPoint(72, 1.0, 1.0, 1.0)
    colorTransferFunction.AddRGBPoint(255, 1.0, 1.0, 1.0)

    # The property describes how the data will look
    volumeProperty = vtk.vtkVolumeProperty()

    # Setup default properties
    volumeProperty.SetColor(colorTransferFunction)
    volumeProperty.SetScalarOpacity(opacityTransferFunction)
    volumeProperty.ShadeOn()
    volumeProperty.SetInterpolationTypeToLinear()

    volumeMapper = vtk.vtkVolumeTextureMapper3D()
    volumeMapper.SetInput(volumeReader.GetOutput())

    volume.SetProperty(volumeProperty)
    volume.SetMapper(volumeMapper)

    #Set up Moving Image 3D renderwindow, with interaction
    ren = vtk.vtkRenderer()
    ren.AddVolume(volume)
    rw = vtk.vtkRenderWindow()
    rw.AddRenderer(ren)
    iren = vtk.vtkRenderWindowInteractor()
    iren.SetRenderWindow(rw)
    rw.SetSize(512,512)

    # This portion of the code clips the mace with the vtkPlanes implicit
    # function.  The clipped region is colored green.
    planes = vtk.vtkPlanes()
    clipper = vtk.vtkClipVolume()
    clipper.SetInput(0,volumeReader.GetOutput())
    clipper.SetClipFunction(planes)
    clipper.InsideOutOn()
    selectMapper = vtk.vtkVolumeTextureMapper3D()
    selectMapper.SetInputConnection(clipper.GetOutputPort())
    selectActor = vtk.vtkVolume()
    selectActor.SetMapper(selectMapper)
    #selectActor.GetProperty().SetColor(0, 1, 0)
    selectActor.VisibilityOff()
    selectActor.SetScale(1.01, 1.01, 1.01)

    # The box widget observes the events invoked by the render window
    # interactor.  These events come from user interaction in the render
    # window.
    boxWidget = vtk.vtkBoxWidget()
    boxWidget.SetInteractor(iren)
    boxWidget.SetPlaceFactor(1.0)

    # As the box widget is interacted with, it produces a transformation
    # matrix that is set on the actor.
    t = vtk.vtkTransform()
    def TransformActor(obj, event):
        global t
        obj.GetTransform(t)
        selectActor.SetUserTransform(t)

    # Place the interactor initially. The actor is used to place and scale
    # the interactor. An observer is added to the box widget to watch for
    # interaction events. This event is captured and used to set the
    # transformation matrix of the actor.
    boxWidget.SetProp3D(selectActor)
    boxWidget.PlaceWidget()
    boxWidget.AddObserver("InteractionEvent", TransformActor)


    style = myVTKRenderWindowInteractorStyle.myVTKRenderWindowInteractorStyle(iren)
    #style = vtk.vtkInteractorStyleTrackballCamera()
    iren.SetInteractorStyle(style)

    rw.Render()
    iren.Start()

if __name__ == "__main__":
    InitializeVolume()